import torch
import numpy as np
from torch import nn
from torch.utils import data
import matplotlib.pyplot as plt
true_w = torch.tensor([2]).float()
true_b = 4
'''生成 y=Xw+b+噪声'''
def synthetic_data(w,b,num_examples):    # num_examples:n个样本
    '''生成 y=Xw+b+噪声'''
    X = torch.normal(0,1,(num_examples,len(w)))  #生成 X，他是一个均值为0，方差为1的随机数，他的大小: 行为num_examples，列为w的长度表示多少个feature
    y = torch.matmul(X,w) + b
    y += torch.normal(0,0.5,y.shape)            #加入一些噪音，均值为0 ，方差为0.01，形状和y是一样
    return X, y.reshape((-1,1))

features, labels = synthetic_data(true_w, true_b, 100)

class linear(torch.nn.Module):
    def __init__(self):
        super(linear, self).__init__()
        self.linear = torch.nn.Linear(1,1)

    def forward(self,x):
        y_pred = self.linear(x)
        return y_pred

model = linear()

loss = torch.nn.MSELoss(reduction='sum')

optimizer = torch.optim.SGD(model.parameters(),lr = 0.001)

epoch_list = []
loss_list = []

for epoch in range(100):
    y_pred = model(features)
    loss_temp = loss(y_pred,labels)
    print(epoch,loss_temp.item())

    optimizer.zero_grad()
    loss_temp.backward()
    optimizer.step()
    epoch_list.append(epoch)
    loss_list.append(loss_temp.item())
print('w = ',model.linear.weight.item())
print('b = ',model.linear.bias.item())
a = model.linear.weight.item()
b = model.linear.bias.item()
x_plot = np.linspace(-5, 5, 5)
y_plot = x_plot * 2 + 4
x_ = np.linspace(-5, 5, 5)
y_ = x_plot * a + b
plt.plot(x_.reshape(-1,1), y_.reshape(-1,1), color='green')
plt.plot(x_plot.reshape(-1,1), y_plot.reshape(-1,1), color='red')
plt.scatter(features, labels)
plt.legend(["model","True function"])
plt.show()


plt.plot(epoch_list, loss_list)
plt.xlabel('times')
plt.ylabel('loss')
plt.title('SGD')
plt.show()
